import numpy
from numpy.linalg import norm
from copy import deepcopy
from time import time
from time import strftime
from time import gmtime
from math import exp
from math import floor
from random import uniform
from random import sample
from random import randint
from random import shuffle
from functools import wraps
import matplotlib.pyplot as plt
from scipy.signal import savgol_filter

plt.rc('font', family='Times New Roman', size=12)


def roulette(fitness: list, rounds: int, criteria='minimum'):
    """
    :param fitness: a list of fitness
    :param rounds: roulette rounds
    :param criteria: minimum or maximum
    :return: indexes of jackpots
    """
    if numpy.array(fitness).all() >= 0:
        pass
    else:
        print('\033[31mError in roulette: fitness < 0. ')
        quit(-1)

    if criteria == 'minimum':
        fitness = [1.0 / i for i in fitness]
    elif criteria == 'maximum':
        fitness = [1.0 / i for i in fitness]
    else:
        print("Error in roulette, preference value {} invalid".format(criteria))
        quit(-1)

    total = sum(fitness)
    if total == 0:
        total = 1

    ratio = [i / total for i in fitness]
    for i in range(1, len(ratio)):
        ratio[i] = ratio[i - 1] + ratio[i]
    jackpots = []
    for t in range(rounds):
        jackpot = -1
        p = uniform(0, 1)
        for i in range(0, len(ratio)):
            if p < ratio[i]:
                jackpot = i
                break
        jackpots.append(jackpot)
    return jackpots


def tournament(candidates: list, p_bts: float, k=2, criteria='minimum'):
    competitors = sample(candidates, k=k)
    if criteria == 'minimum':
        return min(competitors) if uniform(0, 1) < p_bts else max(competitors)
    if criteria == 'maximum':
        return max(competitors) if uniform(0, 1) < p_bts else min(competitors)
    else:
        print(f"Error in tournament, criteria={criteria}")
        exit(-1)


def give_me_some_ints(region: tuple, digits=2):
    """
    :param region: low and high region of random integers
    :param digits: the number of random integers
    :return: digits number of random integers between region[0] and region[1]
    """
    randoms = set()
    randoms.clear()
    while len(randoms) < digits:
        randoms.add(randint(region[0], region[1]))
    randoms = list(randoms)
    randoms.sort()
    return randoms


def _2opt(route: list, distance_func, target=None, depth=100, segment=None, data=None):
    """
    :param route: the optimization target route
    :param distance_func: func to calculate the distance of the route
    :param target: the target times of improvement
    :param depth: stagnation limit
    :param segment: segment of the route to perform 2-opt
    :return: best_route, best_dis, evaluations, improvement, stagnation
    """
    if segment is None:
        segment = (0, len(route) - 1)
    if target is None:
        target = int(1e9)
    best_route = deepcopy(route)
    best_dis = distance_func(best_route)
    evaluations = 1

    # print(best_route)
    improvement, stagnation = 0, 0
    depth_reached = False
    while True:
        # try to break all possible 2 pair of edges by sequence
        # i, j are the break points
        flag = False
        for i in range(segment[0], segment[1] + 1):
            # break point at i
            for j in range(i + 1, segment[1] + 1):
                # break point at j
                new_route = __2opt_swap(best_route, segment=(i, j))
                new_dis = distance_func(new_route)
                evaluations = evaluations + 1
                if best_dis > new_dis:  # better route found, start all over
                    best_dis = new_dis
                    best_route = deepcopy(new_route)
                    flag = True
                    # route improved
                    improvement = improvement + 1
                    stagnation = 0
                    if data is not None:
                        data.draw_route(best_route, dot_size=60)
                    break
                stagnation = stagnation + 1
                if stagnation + 1 > depth:  # stagnation limit reached
                    depth_reached = True
                    flag = True
                    break
            if flag:
                break
        if improvement > target:  # improvement target reached
            break
        if depth_reached:  # stagnation limit reached
            break

    return {'route': best_route,
            'distance': best_dis,
            'evaluations': evaluations,
            'improvement': improvement,
            'stagnation': stagnation}


def __2opt_swap(route: list, segment: tuple):
    _route = deepcopy(route)
    _a = _route[:segment[0]]
    if segment[0] == 0:
        _b = _route[segment[1]::-1]
    else:
        _b = _route[segment[1]:segment[0] - 1:-1]
    _c = _route[segment[1] + 1:]
    _route = _a + _b + _c
    return _route

def two_opt(route, distance_func):
    improve = 0
    for i in range(1, len(route)-1):
        _j, min_diff = -1, 1e99
        for j in range(i+1, len(route)-1):
            diff = distance_func[route[i-1]][route[j]] + distance_func[route[i]][route[j+1]] \
                   - distance_func[route[i-1]][route[i]] - distance_func[route[j]][route[j+1]]
            if min_diff > diff:
                min_diff = diff
                _j = j
        if min_diff < 0 and abs(min_diff) > 1e-6:
            route = route[:i] + route[_j:i-1:-1] + route[_j+1:]
            improve += min_diff
    return route, improve

def similarity(routes1, routes2):
    """
    calculate the similarity of route 1 and routes 2
    :param routes1:
    :param routes2:
    :return:
    """
    len1, len2 = sum(len(route) for route in routes1), sum(len(route) for route in routes2)
    if len1 != len2:  # if route1 and route2 are not equal on length
        print('\033[31mError in func<similarity>, len(route1) != len(route2). ')
        quit(-1)

    res = 0
    for a, b in zip(routes1, routes2):
        res += LCSubstring(a, b)  # count the LCSubstring of a and b
    return float(res) / float(len1)


def LCS(string1, string2):
    """
    Longest Common Subsequence
    :param string1:
    :param string2:
    :return: the length of the longest common subsequence between string1 and string2
    """
    len1 = len(string1)
    len2 = len(string2)
    res = [[0 for i in range(len1 + 1)] for j in range(len2 + 1)]
    for i in range(1, len2 + 1):
        for j in range(1, len1 + 1):
            if string2[i - 1] == string1[j - 1]:
                res[i][j] = res[i - 1][j - 1] + 1
            else:
                res[i][j] = max(res[i - 1][j], res[i][j - 1])
    return res[-1][-1]


def LCSubstring(string1, string2):
    """
    Longest Common Substring
    :param string1:
    :param string2:
    :return: the length of the longest common subsequence between string1 and string2
    """
    len1 = len(string1)
    len2 = len(string2)
    res = [[0 for i in range(len1 + 1)] for j in range(len2 + 1)]
    result = 0
    for i in range(1, len2 + 1):
        for j in range(1, len1 + 1):
            if string2[i - 1] == string1[j - 1]:
                res[i][j] = res[i - 1][j - 1] + 1
                result = max(result, res[i][j])
    return result

class Lifecycle:
    """
        Lifecycle of an optimizer running progress
    """

    def __init__(self, iterations=None, evaluations=None, deadline=None):
        self.iterations = int(1e9) if iterations is None else iterations
        self.evaluations = int(1e9) if evaluations is None else evaluations
        self.deadline = int(1e9) if deadline is None else deadline * 60  # minutes
        self.it = 0
        self.eva = 0
        self.__start_time = time()
        self.elapsed = time() - self.__start_time
        self.__pre = self.elapsed
        self.status = 'mounted'
        self.message = None
        self.gBest = []
        self.data = None

    def set_stopping_condition(self, iterations=None, evaluations=None, deadline=None):
        self.iterations = int(1e9) if iterations is None else iterations
        self.evaluations = int(1e9) if evaluations is None else evaluations
        self.deadline = int(1e9) if deadline is None else deadline * 60  # minutes

    def update(self, it=None, eva=None, timing=True, gBest=None, data=None):
        if it is not None:
            self.it += it
        if eva is not None:
            self.eva += eva
        if timing:
            self.elapsed = time() - self.__start_time
        if gBest is not None:
            self.gBest.append(gBest)
        self.data = data
        self.status = 'alive'

    def check_stopping_condition(self, mute=True):
        msg, stop = ['', '', ''], False
        if self.it >= self.iterations:
            msg[0] = 'iterations limitation reached. '
            stop = True
        if self.eva >= self.evaluations:
            msg[1] = 'evaluations limitation reached. '
            stop = True
        if self.elapsed >= self.deadline:
            msg[2] = 'deadline reached. '
            stop = True
        if stop:
            self.message = msg[0] + msg[1] + msg[2]
            self.status = 'closed'
            t = strftime('%H:%M:%S', gmtime(self.elapsed))
            if mute is False:
                print(f'Terminated at {self.it}, time cost {t}s')

        return stop

    def report(self, per_sec=True, graph=True, mapping=None):
        now = round(self.elapsed)
        if now != self.__pre or not per_sec:
            t = strftime('%H:%M:%S', gmtime(self.elapsed))
            print(f'{t} at {self.it:5d} route length {self.data[-1]:.2f}')
            self.__pre = now
            if graph:
                if mapping is not None:
                    plt.clf()
                    plt.subplot(1, 2, 1)
                    plt.plot(range(len(self.gBest)), self.gBest, linestyle='-', label='gBest')
                    plt.title(label=f'iterations {self.it}, d: {self.gBest[-1]:.2f}')
                    ax = plt.gca()
                    ax.spines['top'].set_color('none')
                    ax.spines['right'].set_color('none')
                    plt.legend()
                    plt.subplot(1, 2, 2)
                    route = mapping[0]
                    draw = mapping[1]
                    draw(plt, route, dot_size=20, show_text=True)
                    ax = plt.gca()
                    # ax.spines['top'].set_color('none')
                    # ax.spines['right'].set_color('none')
                    # ax.spines['left'].set_color('none')
                    # ax.spines['bottom'].set_color('none')
                    ax.set_aspect('equal')
                    plt.pause(0.001)
                else:
                    plt.clf()
                    plt.plot(range(len(self.gBest)), self.gBest, linestyle='-', label='gBest')
                    plt.title(label=f'iterations {self.it}, d: {self.gBest[-1]:.2f}')
                    plt.pause(0.001)

def normalized_vec(vec):
    return vec / norm(vec)

def cake_cutter(n: int, cake: list, fair=True, seq=True):
    """
    divide cake into n parts
    @param n: pieces
    @param cake: cake list
    @param fair: fair cut
    @param seq: keep original sequence
    @return:
    """
    size = len(cake)
    parts = []
    if fair:
        grid, mod = floor(size / n), size % n
        for i in range(n - 1):
            parts.append(grid)
        parts.append(grid + mod)
    else:
        parts = factorization(len(cake), n)
    shuffle(parts)
    res = []
    if seq:
        s = 0
        for p in parts:
            res.append(cake[s:s + p])
            s += p
    else:
        for p in parts:
            sam = sample(cake, k=p)
            res.append(sam)
            for s in sam:
                cake.remove(s)

    return res


def factorization(total: int, n: int):
    """
    @param total: the total number to factorize
    @param n: number of partitions
    @return: a list of n integers summed up to total
    """
    ruler = [i for i in range(0, total)]  # a rule range from 0 to total
    if total < n:
        scales = []
        for i in range(n):
            scales.append(sample(ruler, k=1)[0])
    else:
        scales = sample(ruler, k=n - 1)  # sample n-1 scales on the rule
    scales.sort()  # sort scales
    span = [scales[0] - 0]  # calculate the spans between scales
    for i in range(1, len(scales)):
        span.append(scales[i] - scales[i - 1])
    span.append(total - scales[-1])
    return span


def minimal_insertion(distance_func, candidates, route=None, depot=None):
    if route is None:
        if depot is None:
            print('\033[31mError in func<minimum_insertion>, route=None, depot=None. ')
            quit(-1)
        else:
            route = [depot, depot]
    if depot in candidates:
        candidates.remove(depot)
    while len(candidates) > 0:
        node = sample(candidates, k=1)[0]
        candidates.remove(node)
        _i, min_cost = -1, 1e99
        for i in range(1, len(route)):
            pre = route[i - 1]
            pos = route[i]
            cost = distance_func[pre][node] + distance_func[node][pos] - distance_func[pre][pos]
            if min_cost > cost:
                min_cost = cost
                _i = i
        if _i == -1:
            print('\033[31mError in func<minimum_insertion>, _i == -1. ')
            quit(-1)
        route.insert(_i, node)

    return route

def greedy_sequence(distance_func, candidates, route=None, depot=None):
    if route is None:
        if depot is None:
            print('\033[31mError in func<minimum_insertion>, route=None, depot=None. ')
            quit(-1)
        else:
            route = [depot]
    candidates = candidates[:]
    if depot in candidates:
        candidates.remove(depot)
    while len(candidates) > 0:
        post, min_dis = -1, 1e9
        curr = route[-1]
        for node in candidates:
            dis = distance_func[curr][node]
            if min_dis > dis:
                min_dis = dis
                post = node
        route.append(post)
        candidates.remove(post)

    route.append(depot)
    return route

def edge_set(route: list):
    E = set()
    for i in range(1, len(route)):
        a, b = route[i-1], route[i]
        if a > b:
            a, b = b, a
        E.add((a, b))
    return E


# from optimizer.tsp.tsp_data import TSP
# data = TSP(r'../../data/TSPLIB/burma14.tsp.txt')
# route_mi = minimal_insertion(distance_func=data.distances, candidates=data.nodes[:], depot=1)
# route_gs = greedy_sequence(distance_func=data.distances, candidates=data.nodes[:], depot=1)
#
# route_mi = [1, 2, 1, 3, 1]
# print(route_mi)
# print(edge_set(route_mi))

# for i in range(10000):
#     route_mi = two_opt(route=route_mi, distance_func=data.distances)[0]
#     route_gs = two_opt(route=route_gs, distance_func=data.distances)[0]
#
#     plt.clf()
#     plt.subplot(1, 2, 1)
#     data.draw(plot=plt, route=route_mi, title=f'{data.cost(route=route_mi):.2f}')
#     plt.subplot(1, 2, 2)
#     data.draw(plot=plt, route=route_gs, title=f'{data.cost(route=route_gs):.2f}')
#     plt.pause(0.001)

